import oneflow as torch
import oneflow.nn as nn
import oneflow.nn.functional as F


class TimeReductionLayer(nn.Module):
    def __init__(self, ntimes=2):
        super(TimeReductionLayer, self).__init__()

        self.ntimes = ntimes

    def forward(self, inputs, inputs_length):

        batch_size, feat_len, feat_dim = inputs.size()
        if feat_len % self.ntimes != 0:
            pad_len = self.ntimes - feat_len % self.ntimes
            inputs = F.pad(inputs, pad=(0, 0, 0, pad_len), value=0.0)
        else:
            pad_len = 0

        inputs = inputs.reshape(batch_size, -1, feat_dim * 2)
        inputs_length = torch.ceil(inputs_length.float() / 2).long()
        
        return inputs, inputs_length